from fractions import Fraction
from inspect import cleandoc
import time
import requests
from PIL import Image, ImageSequence, ImageOps
from io import BytesIO
import torch
import numpy as np
import av


class LoadRemoteImage:
    """
    一个从URL加载图像的节点。

    类方法
    -------------
    INPUT_TYPES (dict):
        告诉主程序节点的输入参数。
    IS_CHANGED:
        可选方法，用于控制节点何时重新执行。

    属性
    ----------
    RETURN_TYPES (`tuple`):
        输出元组中每个元素的类型。
    RETURN_NAMES (`tuple`):
        可选：输出元组中每个输出的名称。
    FUNCTION (`str`):
        入口点方法的名称。例如，如果 `FUNCTION = "execute"`，则它将运行 Example().execute()
    OUTPUT_NODE ([`bool`]):
        如果此节点是输出节点，则会从图表中输出结果/图像。SaveImage节点就是一个例子。
        后端会迭代这些输出节点，并尝试执行它们的所有父节点（如果它们的父图连接正确）。
        如果不存在，则默认为False。
    CATEGORY (`str`):
        节点应在UI中显示的类别。
    execute(s) -> tuple || None:
        入口点方法。此方法的名称必须与`FUNCTION`属性的值相同。
        例如，如果 `FUNCTION = "execute"`，则此方法的名称必须是 `execute`，如果 `FUNCTION = "foo"`，则必须是 `foo`。
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        """
        返回一个包含所有输入字段配置的字典。
        一些类型（字符串）："MODEL"、"VAE"、"CLIP"、"CONDITIONING"、"LATENT"、"IMAGE"、"INT"、"STRING"、"FLOAT"。
        输入类型 "INT"、"STRING" 或 "FLOAT" 是节点上字段的特殊值。
        类型可以是用于选择的列表。

        返回: `dict`:
            - 键 input_fields_group (`string`): 可以是 required、hidden 或 optional。节点类必须有 `required` 属性
            - 值 input_fields (`dict`): 包含输入字段配置：
                * 键 field_name (`string`): 入口点方法参数的名称
                * 值 field_config (`tuple`):
                    + 第一个值是指示字段类型的字符串或用于选择的列表。
                    + 第二个值是类型 "INT"、"STRING" 或 "FLOAT" 的配置。
        """
        return {
            "required": {
                "image_url": ("STRING", {"multiline": True}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    DESCRIPTION = cleandoc(__doc__)
    FUNCTION = "load_image"

    # OUTPUT_NODE = False
    # OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node

    CATEGORY = "image"

    def pil_to_tensor(self, pil_image):
        """
        将PIL图像（或图像序列）转换为torch张量。

        参数:
            pil_image (PIL.Image): 输入图像，可能是多帧（例如GIF）。

        返回:
            tuple[torch.Tensor, torch.Tensor]:
                - image_tensor: 浮点张量，范围[0,1]，形状 (frames, H, W, 3)
                - mask_tensor:  浮点张量，范围[0,1]，形状 (frames, H, W)
        """
        image_tensors = []
        mask_tensors = []

        for frame in ImageSequence.Iterator(pil_image):
            # Handle EXIF orientation
            frame = ImageOps.exif_transpose(frame)

            # Normalize 32-bit integer images
            if frame.mode == "I":
                frame = frame.point(lambda p: p * (1 / 255))

            # Convert to RGB float tensor
            rgb_array = np.array(frame.convert("RGB"), dtype=np.float32) / 255.0
            rgb_tensor = torch.from_numpy(rgb_array).unsqueeze(0)  # (1, H, W, 3)

            # Extract alpha channel as mask if available
            if "A" in frame.getbands():
                alpha_array = np.array(frame.getchannel("A"), dtype=np.float32) / 255.0
                mask_tensor = 1.0 - torch.from_numpy(alpha_array)  # invert alpha
            else:
                # Fallback: dummy mask (64x64, all zeros)
                mask_tensor = torch.zeros((64, 64), dtype=torch.float32)

            image_tensors.append(rgb_tensor)
            mask_tensors.append(mask_tensor.unsqueeze(0))  # add frame dimension

        # Concatenate frames if multiple
        if len(image_tensors) > 1:
            image_tensor = torch.cat(image_tensors, dim=0)
            mask_tensor = torch.cat(mask_tensors, dim=0)
        else:
            image_tensor = image_tensors[0]
            mask_tensor = mask_tensors[0]

        return image_tensor, mask_tensor

    def load_image(self, image_url):
        # download image
        response = requests.get(image_url)
        img = Image.open(BytesIO(response.content))
        return self.pil_to_tensor(img)

    """
        如果任何输入发生变化，节点将始终重新执行，但
        即使输入没有变化，也可以使用此方法强制节点再次执行。
        您可以让此节点返回一个数字或字符串。这个值将与节点上次执行时返回的值进行比较，
        如果不同，节点将再次执行。
        此方法在核心仓库的LoadImage节点中使用，其中它们将图像哈希作为字符串返回，如果图像哈希
        在执行之间发生变化，则LoadImage节点将再次执行。
    """

    @classmethod
    def IS_CHANGED(cls, image_url):
        return image_url


class LoadRemoteVideo:
    """
    一个从URL加载视频的节点。

    类方法
    -------------
    INPUT_TYPES (dict):
        定义节点所需的输入参数。
    IS_CHANGED:
        可选方法，用于控制节点何时重新执行。

    属性
    ----------
    RETURN_TYPES (`tuple`):
        输出元组中每个元素的类型。
    FUNCTION (`str`):
        要运行的入口点方法。
    CATEGORY (`str`):
        节点应在UI中显示的类别。
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "video_url": ("STRING", {"multiline": True}),
                "start_second": ("INT", {"multiline": False}),
                "video_length_seconds": ("INT", {"multiline": False}),
            },
        }

    RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT")
    RETURN_NAMES = ("IMAGES", "AUDIO", "FPS")
    DESCRIPTION = cleandoc(__doc__)
    FUNCTION = "load_video"
    CATEGORY = "image/video"

    def load_video(self, video_url, start_second, video_length_seconds):
        """
        从URL下载视频并在单个循环中将请求的片段解码为张量
        （包括视频和音频）。
        """

        # download to buffer
        download_start = time.time()
        response = requests.get(video_url, stream=True)
        response.raise_for_status()
        buffer = BytesIO(response.content)
        download_end = time.time()
        print(f"Finished downloading in {download_end - download_start:.2f}s.")

        # open container
        container = av.open(buffer)

        # pick streams
        video_stream = next(s for s in container.streams if s.type == "video")
        audio_stream = next((s for s in container.streams if s.type == "audio"), None)
        frame_rate = video_stream.average_rate or Fraction(25, 1)

        # compute bounds
        start_pts = int(start_second / av.time_base)  # in microseconds
        end_second = start_second + video_length_seconds

        # seek close to start
        container.seek(start_pts, any_frame=False, backward=True, stream=video_stream)

        frames = []
        audio_frames = []

        # decode interleaved
        for packet in container.demux((video_stream, audio_stream) if audio_stream else (video_stream,)):
            for frame in packet.decode():
                t = float(frame.pts * frame.time_base)
                if t < start_second:
                    continue
                if t >= end_second:
                    container.close()
                    images_tensor = torch.cat(frames, dim=0) if frames else None
                    audio_tensor = None
                    if audio_frames:
                        waveform = torch.cat(audio_frames, dim=1).unsqueeze(0)
                        audio_tensor = {
                            "waveform": waveform,
                            "sample_rate": audio_stream.rate,
                        }
                    return images_tensor, audio_tensor, frame_rate

                if packet.stream.type == "video":
                    img = frame.to_ndarray(format="rgb24")
                    img = torch.from_numpy(img).float() / 255.0
                    frames.append(img.unsqueeze(0))

                elif packet.stream.type == "audio":
                    arr = frame.to_ndarray()  # (channels, samples)
                    audio_frames.append(torch.from_numpy(arr))

        # finalize
        images_tensor = torch.cat(frames, dim=0) if frames else None
        audio_tensor = None
        if audio_frames:
            waveform = torch.cat(audio_frames, dim=1).unsqueeze(0)
            audio_tensor = {"waveform": waveform, "sample_rate": audio_stream.rate}

        return images_tensor, audio_tensor, frame_rate

    @classmethod
    def IS_CHANGED(cls, video_url, start_second, video_length_seconds):
        return video_url


class LoadRemoteAudio:
    """
    一个示例节点

    类方法
    -------------
    INPUT_TYPES (dict):
        告诉主程序节点的输入参数。
    IS_CHANGED:
        可选方法，用于控制节点何时重新执行。

    属性
    ----------
    RETURN_TYPES (`tuple`):
        输出元组中每个元素的类型。
    RETURN_NAMES (`tuple`):
        可选：输出元组中每个输出的名称。
    FUNCTION (`str`):
        入口点方法的名称。例如，如果 `FUNCTION = "execute"`，则它将运行 Example().execute()
    OUTPUT_NODE ([`bool`]):
        如果此节点是输出节点，则会从图表中输出结果/图像。SaveImage节点就是一个例子。
        后端会迭代这些输出节点，并尝试执行它们的所有父节点（如果它们的父图连接正确）。
        如果不存在，则默认为False。
    CATEGORY (`str`):
        节点应在UI中显示的类别。
    execute(s) -> tuple || None:
        入口点方法。此方法的名称必须与`FUNCTION`属性的值相同。
        例如，如果 `FUNCTION = "execute"`，则此方法的名称必须是 `execute`，如果 `FUNCTION = "foo"`，则必须是 `foo`。
    """

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        """
        返回一个包含所有输入字段配置的字典。
        一些类型（字符串）："MODEL"、"VAE"、"CLIP"、"CONDITIONING"、"LATENT"、"IMAGE"、"INT"、"STRING"、"FLOAT"。
        输入类型 "INT"、"STRING" 或 "FLOAT" 是节点上字段的特殊值。
        类型可以是用于选择的列表。

        返回: `dict`:
            - 键 input_fields_group (`string`): 可以是 required、hidden 或 optional。节点类必须有 `required` 属性
            - 值 input_fields (`dict`): 包含输入字段配置：
                * 键 field_name (`string`): 入口点方法参数的名称
                * 值 field_config (`tuple`):
                    + 第一个值是指示字段类型的字符串或用于选择的列表。
                    + 第二个值是类型 "INT"、"STRING" 或 "FLOAT" 的配置。
        """
        return {
            "required": {
                "audio_url": ("STRING", {"multiline": True}),
            },
        }

    RETURN_TYPES = ("AUDIO",)
    RETURN_NAMES = ("AUDIO",)
    DESCRIPTION = cleandoc(__doc__)
    FUNCTION = "load_audio"
    CATEGORY = "audio"

    def load_audio(self, audio_url):
        """
        从URL下载音频并将请求的片段解码为张量。
        """

        # download to buffer
        download_start = time.time()
        response = requests.get(audio_url, stream=True)
        response.raise_for_status()
        buffer = BytesIO(response.content)
        download_end = time.time()
        print(f"Finished downloading in {download_end - download_start:.2f}s.")

        # open container
        container = av.open(buffer)

        # pick audio stream
        audio_stream = next((s for s in container.streams if s.type == "audio"), None)
        if not audio_stream:
            raise ValueError("No audio stream found in the file.")

        audio_frames = []

        # decode interleaved
        for packet in container.demux(audio_stream):
            for frame in packet.decode():
                arr = frame.to_ndarray()  # (channels, samples)
                audio_frames.append(torch.from_numpy(arr))

        print(f"Decoded {len(audio_frames)} audio frames.")

        waveform = torch.cat(audio_frames, dim=1).unsqueeze(0)
        audio_tensor = {"waveform": waveform, "sample_rate": audio_stream.rate}

        print(f"Final audio tensor shape: {audio_tensor['waveform'].shape}, sample rate: {audio_tensor['sample_rate']}")

        return (audio_tensor,)

    """
        如果任何输入发生变化，节点将始终重新执行，但
        即使输入没有变化，也可以使用此方法强制节点再次执行。
        您可以让此节点返回一个数字或字符串。这个值将与节点上次执行时返回的值进行比较，
        如果不同，节点将再次执行。
        此方法在核心仓库的LoadImage节点中使用，其中它们将图像哈希作为字符串返回，如果图像哈希
        在执行之间发生变化，则LoadImage节点将再次执行。
    """

    @classmethod
    def IS_CHANGED(s, audio_url):
        return audio_url
